package requester

import (
	"bytes"
	"compress/gzip"
	"context"
	"crypto/tls"
	"io"
	"mime/multipart"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/andybalholm/brotli"
	"golang.org/x/net/html/charset"
	"golang.org/x/net/proxy"
	"golang.org/x/net/publicsuffix"
	"golang.org/x/text/encoding/unicode"
	"golang.org/x/text/transform"

	"gitee.com/sillyman/WebCrawler/urlqueue"
)

type (
	// RequestHook 在发送请求前调用，返回 false 则表示此消息不发出
	// 一般用于修改 request 消息，或者阻止发送，或打印请求日志
	RequestHook func(*http.Request) bool

	// ResponseHook 在收到应答消息后调用
	// 用于分析应答的HTML页面：下载资源、抽取指定的URLs加入列队
	ResponseHook func(URL *url.URL, statusCode int, contentType string, body io.Reader, header http.Header, cookies []*http.Cookie)
)

// Requester HTTP请求者
type Requester struct {
	// requestHook 发送请求前调用
	requestHook RequestHook

	// responseHook 接收到应答后调用
	responseHook ResponseHook

	// CustomHeader 额外添加的自定义头
	CustomHeader http.Header

	// CustomCookies 额外添加的cookies
	CustomCookies []*http.Cookie

	// client 是HTTP客户端
	client *http.Client

	// clientJar cookies管理器
	clientJar *cookiejar.Jar

	// clientTransport http传输层
	clientTransport *http.Transport
}

// NewRequester 创建一个HTTP请求者实例
// 忽略HTTPs验证的过程，支持gzip,deflate,br内容，自动将非UTF8编码的页面转换成UTF8格式
func NewRequester(timeout time.Duration, onReqFunc RequestHook, onResFunc ResponseHook) *Requester {
	r := &Requester{
		requestHook:  onReqFunc,
		responseHook: onResFunc,
		CustomHeader: http.Header(map[string][]string{
			"Accept-Encoding": {"gzip", "deflate", "br"}, // br(Brotli) 是一种全新的数据压缩格式，2015年发布
			"User-Agent":      {"Mozilla/5.0 (X11; Linux x86_64; rv:83.0) Gecko/20100101 Firefox/83.0"},
		}),
	}

	r.clientJar, _ = cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	r.clientTransport = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	r.client = &http.Client{
		Timeout:   timeout,
		Jar:       r.clientJar,
		Transport: r.clientTransport,
	}
	return r
}

// SetRequestHook 设置请求钩子
func (r *Requester) SetRequestHook(onReqFunc RequestHook) { r.requestHook = onReqFunc }

// SetResponseHook 设置应答钩子
func (r *Requester) SetResponseHook(onResFunc ResponseHook) { r.responseHook = onResFunc }

// SetProxy 设置代理，格式：`http[s]://username:pwd@host:port`
func (r *Requester) SetProxy(rawurl string) error {
	URL, err := url.Parse(rawurl)
	if err != nil {
		return err
	}
	r.clientTransport.Proxy = http.ProxyURL(URL)
	return nil
}

// SetProxyBySocks5 设置 Socks5 代理
func (r *Requester) SetProxyBySocks5(addr, username, password string) error {
	dialer, err := proxy.SOCKS5("tcp", addr,
		&proxy.Auth{User: username, Password: password},
		proxy.Direct,
	)
	if err != nil {
		return err
	}

	dc := dialer.(interface {
		DialContext(ctx context.Context, network, addr string) (net.Conn, error)
	})
	r.clientTransport.DialContext = dc.DialContext
	return nil
}

// request 发送请求，获取应答消息
// 对于应答消息被 gzip 或 br 压缩，将自动解压缩。
// 如果是 HTML 页面，如果字符编码不是 UTF-8，将自动转换成 UTF-8
func (r *Requester) Request(method, contentType string, myURL *urlqueue.MyURL, body io.Reader) error {
	req, err := http.NewRequest(method, myURL.String(), body)
	if err != nil {
		return err
	}

	// 添加额外的 cookie 和 header
	{
		for _, cookie := range r.CustomCookies {
			req.AddCookie(cookie)
		}
		for key, values := range r.CustomHeader {
			for _, v := range values {
				req.Header.Add(key, v)
			}
		}
	}

	// 设置请求头的 Content-Type 和 Referer
	{
		if contentType != "" {
			req.Header.Set("Content-Type", contentType)
		}
		if myURL.Referer != "" {
			req.Header.Set("Referer", myURL.Referer)
		}
	}

	// 先执行请求前的钩子函数，然后发出请求
	if r.requestHook != nil && !r.requestHook(req) {
		return nil
	}
	res, err := r.client.Do(req)
	if err != nil {
		return err
	}
	defer func() {
		if res.Body != nil {
			res.Body.Close()
		}
	}()

	if strings.ToUpper(req.Method) != http.MethodHead {
		resContentType := res.Header.Get("Content-Type")

		// 对于应答消息，也行添加解压器和编码转换器
		reader, err := r.maybeAddDecompressor(res.Header.Get("Content-Encoding"), res.Body)
		if err != nil {
			return err
		}
		if strings.Contains(strings.ToLower(resContentType), "text/html") { // 仅仅尝试对HTML页面进行编码转换
			reader, err = r.maybeAddEncodingDecoder(resContentType, reader)
			if err != nil {
				return err
			}
		}

		// 执行 HTTP 应答钩子函数
		if r.responseHook != nil {
			r.responseHook(myURL.URL, res.StatusCode, resContentType, reader, res.Header, res.Cookies())
		}
	}

	return nil
}

// Get 发送 GET 请求
func (r *Requester) Get(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodGet, "", myURL, strings.NewReader(values.Encode()))
}

// Post 发送 POST 请求
// 普通表单提交，content-type=application/x-www-form-urlencoded
func (r *Requester) Post(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodPost, "application/x-www-form-urlencoded", myURL, strings.NewReader(values.Encode()))
}

// PostCustomContentType 发送 POST 请求，支持定义 ContentType 头
// 如果向服务器发送 JSON 内容时，应该设置为 application/json
func (r *Requester) PostCustomContentType(myURL *urlqueue.MyURL, contentType string, body io.Reader) error {
	if contentType == "" {
		contentType = "application/x-www-form-urlencoded"
	}
	return r.Request(http.MethodPost, contentType, myURL, body)
}

// PostUploadFiles 发送 POST 请求上传文件
// 参数：
// fileFieldName 文件框的字段名称
// files 是文件路径
// values 附加的表单参数
func (r *Requester) PostUploadFiles(myURL *urlqueue.MyURL, values url.Values, fileFieldName string, files ...string) error {
	buf := new(bytes.Buffer)
	partWriter := multipart.NewWriter(buf)

	// 添加上传的文件
	for _, pth := range files {
		f, err := os.Open(pth)
		if err != nil {
			return err
		}

		if part, err := partWriter.CreateFormFile(fileFieldName, filepath.Base(pth)); err != nil {
			return err
		} else {
			if _, err := io.Copy(part, f); err != nil {
				return err
			}
		}

		_ = f.Close()
	}

	// 添加表单的其它字段
	for filedName, values := range values {
		for _, v := range values {
			_ = partWriter.WriteField(filedName, v)
		}
	}

	if err := partWriter.Close(); err != nil {
		return err
	}
	return r.Request(http.MethodPost, partWriter.FormDataContentType(), myURL, buf)
}

// Put 发送 PUT 请求
// 普通表单提交，content-type=application/x-www-form-urlencoded
func (r *Requester) Put(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodPut, "application/x-www-form-urlencoded", myURL, strings.NewReader(values.Encode()))
}

// PutCustomContentType 发送 PUT 请求，支持定义 ContentType 头
// 如果向服务器发送 JSON 内容时，应该设置为 application/json
func (r *Requester) PutCustomContentType(myURL *urlqueue.MyURL, contentType string, body io.Reader) error {
	if contentType == "" {
		contentType = "application/x-www-form-urlencoded"
	}
	return r.Request(http.MethodPut, contentType, myURL, body)
}

// Patch 发送 PATCH 请求
// 普通表单提交，content-type=application/x-www-form-urlencoded
func (r *Requester) Patch(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodPatch, "application/x-www-form-urlencoded", myURL, strings.NewReader(values.Encode()))
}

// PatchCustomContentType 发送 PATCH 请求，支持定义 ContentType 头
// 如果向服务器发送 JSON 内容时，应该设置为 application/json
func (r *Requester) PatchCustomContentType(myURL *urlqueue.MyURL, contentType string, body io.Reader) error {
	if contentType == "" {
		contentType = "application/x-www-form-urlencoded"
	}
	return r.Request(http.MethodPatch, contentType, myURL, body)
}

// Delete 发送 DELETE 请求
// 普通表单提交，content-type=application/x-www-form-urlencoded
func (r *Requester) Delete(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodDelete, "application/x-www-form-urlencoded", myURL, strings.NewReader(values.Encode()))
}

// DeleteCustomContentType 发送 DELETE 请求，支持定义 ContentType 头
// 如果向服务器发送 JSON 内容时，应该设置为 application/json
func (r *Requester) DeleteCustomContentType(myURL *urlqueue.MyURL, contentType string, body io.Reader) error {
	if contentType == "" {
		contentType = "application/x-www-form-urlencoded"
	}
	return r.Request(http.MethodDelete, contentType, myURL, body)
}

// Head 发送 HEAD 请求
func (r *Requester) Head(myURL *urlqueue.MyURL, values url.Values) error {
	return r.Request(http.MethodHead, "", myURL, strings.NewReader(values.Encode()))
}

// maybeAddDecompressor 也许会给 body 添加相关的解压器，当`contentEncoding`为 gzip 或 br
func (r *Requester) maybeAddDecompressor(contentEncoding string, body io.Reader) (io.Reader, error) {
	switch strings.ToLower(contentEncoding) {
	case "gzip":
		return gzip.NewReader(body)
	case "br":
		return brotli.NewReader(body), nil
	default:
		return body, nil
	}
}

// determineEncoding 也许会给 body 添加相关的字符编码转换器，当内容不是UTF-8时，会将它转换成UTF-8
// 先读取 body 前 1024 bytes内容，然后由 `charset.DetermineEncoding(data, contentType)` 确认编码，再决定是否进行编码转换
func (r *Requester) maybeAddEncodingDecoder(contentType string, body io.Reader) (io.Reader, error) {
	buf := make([]byte, 1024)
	n, err := body.Read(buf)
	if err != nil {
		return nil, err
	}

	// 确定字符编码
	encode, _, _ := charset.DetermineEncoding(buf[:n], contentType)
	if encode != unicode.UTF8 { // 对于非 UTF-8 编码的页面，则转为 UTF-8
		readers := io.MultiReader(bytes.NewReader(buf[:n]), body)
		return transform.NewReader(readers, encode.NewDecoder()), nil
	}

	return body, nil
}
